package net.aias.springaichat.adapter.ai.config;

import lombok.extern.slf4j.Slf4j;
import net.aias.springaichat.adapter.ai.client.DoubaoEmbeddingClient;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.web.reactive.function.BodyInserters;
import org.springframework.web.reactive.function.client.ClientRequest;
import org.springframework.web.reactive.function.client.ClientResponse;
import org.springframework.web.reactive.function.client.ExchangeFilterFunction;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.nio.charset.StandardCharsets;

/**
 * @Auther: cheng.tang
 * @Date: 2025/5/13
 * @Description: spring-ai-chat
 */
@Configuration
@Slf4j
public class EmbeddingConfig {

    @Bean
    public DoubaoEmbeddingClient doubaoEmbeddingClient() {
//        log.info("DoubaoEmbeddingClient init {} ", System.getenv("DOUBAO_API_KEY"));
        WebClient webClient = WebClient.builder()
                .baseUrl("https://ark.cn-beijing.volces.com/api/v3/embeddings")
                .defaultHeader(HttpHeaders.CONTENT_TYPE, MediaType.APPLICATION_JSON_VALUE)
                .defaultHeader(HttpHeaders.AUTHORIZATION, "Bearer " + System.getenv("DOUBAO_API_KEY"))
//                .filter(logRequest())
                .filter(logResponse())
                .build();
        return new DoubaoEmbeddingClient(webClient);
    }

//
//    private ExchangeFilterFunction logRequest() {
//        return ExchangeFilterFunction.ofRequestProcessor(clientRequest -> {
//            StringBuilder sb = new StringBuilder("\n\n================ API 请求 ================\n");
//            sb.append("URI         : ").append(clientRequest.url()).append("\n");
//            sb.append("Method      : ").append(clientRequest.method()).append("\n");
//            sb.append("Headers     : ").append(clientRequest.headers()).append("\n");
//
//            if (!isJsonMediaType(clientRequest.headers().getContentType())) {
//                sb.append("Request Body: [Non-JSON content]\n");
//                log.info("{}", sb);
//                return Mono.just(clientRequest);
//            }
//
//            // 使用DataBufferUtils读取请求体
//            return readRequestBody(clientRequest)
//                    .flatMap(body -> {
//                        sb.append("Request Body: ").append(body).append("\n");
//                        sb.append("==========================================\n\n");
//                        log.info("{}", sb);
//                        // 重新创建请求
//                        return Mono.just(ClientRequest.from(clientRequest)
//                                .body(BodyInserters.fromValue(body)).build());
//                    });
//        });
//    }

    private ExchangeFilterFunction logResponse() {
        return ExchangeFilterFunction.ofResponseProcessor(clientResponse -> {
            StringBuilder sb = new StringBuilder("================ API 响应 ================\n");
            sb.append("Status Code : ").append(clientResponse.statusCode()).append("\n");
            sb.append("Headers     : ").append(clientResponse.headers().asHttpHeaders()).append("\n");

            if (clientResponse.headers().contentType().isEmpty() ||
                    !isJsonMediaType(clientResponse.headers().contentType().get())) {
                sb.append("Response Body: [Non-JSON content]\n");
                log.info("{}", sb);
//                System.out.println(sb);
                return Mono.just(clientResponse);
            }

            // 使用DataBufferUtils读取响应体
            ClientResponse.Headers headers = clientResponse.headers();
            return clientResponse.bodyToFlux(DataBuffer.class)
                    .collectList()
                    .flatMap(dataBuffers -> {
                        String body = dataBuffers.stream()
                                .map(buffer -> {
                                    byte[] bytes = new byte[buffer.readableByteCount()];
                                    buffer.read(bytes);
                                    DataBufferUtils.release(buffer);
                                    return new String(bytes, StandardCharsets.UTF_8);
                                })
                                .reduce("", (a, b) -> a + b);

                        sb.append("Response Body: ").append(body).append("\n");
                        sb.append("==========================================\n\n");
//                        System.out.println(sb);
                        log.info("{}", sb);
                        // 重新创建响应
                        return Mono.just(ClientResponse.from(clientResponse)
                                .headers(h -> h.putAll(headers.asHttpHeaders()))
                                .body(body) // 改用 bodyValue 直接设置响应体
                                .build());
                    });
        });
    }

    // 读取请求体的辅助方法
    private Mono<String> readRequestBody(ClientRequest request) {
        // 创建一个虚拟的ClientResponse来读取请求体
        return Mono.defer(() -> {
            Flux<DataBuffer> body = request.body() instanceof BodyInserters.MultipartInserter
                    ? Flux.empty() // 不处理multipart
                    : null;

            if (body == null) {
                return Mono.just("");
            }

            return body.collectList()
                    .map(dataBuffers -> {
                        StringBuilder builder = new StringBuilder();
                        for (DataBuffer buffer : dataBuffers) {
                            byte[] bytes = new byte[buffer.readableByteCount()];
                            buffer.read(bytes);
                            DataBufferUtils.release(buffer);
                            builder.append(new String(bytes, StandardCharsets.UTF_8));
                        }
                        return builder.toString();
                    })
                    .defaultIfEmpty("");
        });
    }

    // 判断是否为JSON类型的MediaType
    private boolean isJsonMediaType(MediaType mediaType) {
        return mediaType != null && (
                MediaType.APPLICATION_JSON.isCompatibleWith(mediaType) ||
                        ("application".equals(mediaType.getType()) &&
                                mediaType.getSubtype().endsWith("+json")));
    }


}
